import torch
import torchvision
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
import time

train_data = torchvision.datasets.CIFAR10(root='./dataset', train=True, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)

test_data = torchvision.datasets.CIFAR10(root='./dataset', train=False, download=True, transform=transforms.ToTensor())
test_loader = DataLoader(test_data, batch_size=128, shuffle=True)

print(f"训练集的长度为：{len(train_data)}")
print(f"测试集的长度为：{len(test_data)}")
print(f"测试集的 dataloader 长度为：{len(test_loader)}")

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.model1 = Sequential(
            Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=True, padding=0),
            Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=True, padding=0),
            Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=True, padding=0),
            Flatten(),
            nn.Linear(in_features=64 * 4 * 4, out_features=64),
            nn.Linear(in_features=64, out_features=10)
        )


    def forward(self, x):
        x = self.model1(x)
        return x

model = Model()
print(model)

loss_fn = nn.CrossEntropyLoss()
learning_rate = 0.01
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
epochs = 10
train_step = 0

# tensorboard记录一下训练时的 loss 变换
writer = SummaryWriter("./logs-train")
for epoch in range(epochs):
    print(f"-------------第{epoch}轮训练开始了------------")
    start = time.time()
    model.train()
    for data in train_loader:
        images, labels = data
        out_labels = model(images)
        loss = loss_fn(out_labels, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_step += 1
        if train_step % 100 == 0:
            end = time.time()
            print(f"训练 100次的时间为{end - start}")
            print(f'-------------训练次数：{train_step}，loss 为: {loss.item()}------------')
            writer.add_scalar('train_loss', loss.item(), train_step)

    model.eval()
    test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            out_labels = model(images)
            test_loss += loss_fn(out_labels, labels)
            total_accuracy += (out_labels.argmax(dim=1) == labels).sum()

    writer.add_scalar('test_loss', test_loss / len(test_data), epoch)
    print(f"------------测试集上的 loss: {test_loss / len(test_data)}, accuracy: {total_accuracy / len(test_data)}")
    writer.add_scalar('test_accuracy', total_accuracy / len(test_data), epoch)
    print(f"------------测试集上的 accuracy: {total_accuracy / len(test_data)}")



